import PIL.Image as Image
import numpy as np
from matplotlib import pyplot as plt
from torchvision.utils import make_grid
import networkx as nx
import cv2
import copy
import torch
import sys
import os
import datetime
import pathlib as path
import argparse

from saliency_map import fig2data

sys.path.append("./Graph2plan/Network/")
from model.utils import int_tuple, str_tuple, bool_flag, get_color_map, get_vocab
from model.box_utils import box_rel2abs
from train import (
    check_manual_seed,
    get_model,
    get_dataset,
    get_dataloader,
    get_data_loaders,
    get_optimizer,
    get_scheduler,
    get_losses,
    batch_cuda
)
def get_args():
    parser = argparse.ArgumentParser()

    ''' Dataset '''
    parser.add_argument('--dataset_dir', default='./Graph2plan/Network/data', type=str)
    parser.add_argument('--image_size', default='128,128', type=int_tuple)
    parser.add_argument('--input_dim', default=3, type=int)
    parser.add_argument('--with_house', default='0', type=bool_flag)
    parser.add_argument('--pos_dim', default=25, type=int)
    parser.add_argument('--area_dim', default=10, type=int)

    ''' Dataloader '''
    parser.add_argument('--batch_size', default=20, type=int)
    parser.add_argument('--workers', default=8, type=int)
    parser.add_argument('--train_shuffle', default='1', type=bool_flag)

    ''' Model '''
    # architecture
    parser.add_argument('--gene_layout', default='1', type=bool_flag)
    parser.add_argument('--box_refine', default='1', type=bool_flag)
    # input
    parser.add_argument('--embedding_dim', default=128,type=int)
    # refine
    parser.add_argument('--refinement_dims', default='1024, 512, 256, 128, 64',type=int_tuple)
    # box refine
    parser.add_argument('--box_refine_arch', default='I15,C3-64-2,C3-128-2,C3-256-2',type=str)
    parser.add_argument('--roi_cat_feature',default='1',type=bool_flag)  
    # control
    parser.add_argument('--gt_box', default=0, type=bool_flag)
    parser.add_argument('--relative', default=1, type=bool_flag)

    ''' Loss '''
    parser.add_argument('--mutex', default=1, type=bool_flag)
    parser.add_argument('--inside', default=1, type=bool_flag)
    parser.add_argument('--coverage', default=1, type=bool_flag)
    parser.add_argument('--render', default=1, type=bool_flag)
    parser.add_argument('--nsample', default=100,type=int)
    parser.add_argument('--loss_refine', default=0, type=bool_flag)
    parser.add_argument('--render_refine', default=0, type=bool_flag)    

    ''' Optimizer '''
    parser.add_argument('--optimizer',default='Adam',type=str)
    parser.add_argument('--scheduler',default='plateau',type=str)
    parser.add_argument('--learning_rate', default=1e-4, type=float)
    parser.add_argument('--decay_rate', default=1e-4, type=float)
    parser.add_argument('--step_size', default=10, type=float)
    parser.add_argument('--step_rate', default=0.5, type=float)

    ''' Checkpoints '''
    parser.add_argument('--save_interval', default=5, type=int)
    parser.add_argument('--n_saved', default=20, type=int)
    parser.add_argument('--pretrain', default=None, type=str)
    parser.add_argument('--skip_train', default=0, type=bool_flag)

    ''' Trainer '''
    parser.add_argument('--seed', default=74269,type=int)
    parser.add_argument('--epoch', default=101,type=int)
    parser.add_argument('--start_epoch',default=None,type=int)

    ''' Others '''
    parser.add_argument('--gpu', default='0', type=str)
    parser.add_argument('--multi_gpu', default=None, type=str)
    parser.add_argument('--suffix',default=None,type=str)
    parser.add_argument('--debug', default=0, type=bool_flag)
    parser.add_argument('--test', default=0, type=bool_flag)

    args = parser.parse_args(args=[
        '--skip_train', '1','--pretrain',
        './Graph2plan/experiment/2022-01-26/DeepLayout_2022-01-26_11-36-22/checkpoints/latest_model_101.pth'])
    return args

def load_checkpoint(args):
    args.epoch=args.epoch if not args.debug else 6
    print("Create dir...")
    start_date = str(datetime.datetime.now().strftime('%Y-%m-%d'))+("" if not args.debug else "_debug")+("" if not args.test else "_test")
    start_date = "2022-01-26" # for debug
    if not os.path.exists(f'../experiment'):
        os.mkdir(f'../experiment')
    experiment_dir = path.Path(f'../experiment/{start_date}')
    experiment_dir.mkdir(exist_ok=True)
    start_time = str(datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')) + '' if args.suffix is None else args.suffix
    start_time = "2022-01-26_11-36-22" # for debug
    file_dir = path.Path(f'{experiment_dir}/DeepLayout_{start_time}')
    file_dir.mkdir(exist_ok=True)
    checkpoints_dir = file_dir.joinpath('checkpoints/')
    checkpoints_dir.mkdir(exist_ok=True)

    # check_manual_seed(args)

    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu if args.multi_gpu is None else args.multi_gpu
    print("gpu_availabel:", os.environ["CUDA_VISIBLE_DEVICES"])
    print("Create dataloader...")
    train_loader,valid_loader,test_loader = get_data_loaders(args)
    print("Create model...")
    model = get_model(args)
    model.cuda()
    print("Gene:",model.refinement_net!=None and args.gene_layout)
    print("Refine:",args.box_refine)
    print("Cat feat:",args.roi_cat_feature)
    print("GT BOX:",args.gt_box)
    print("Iniside Loss:",args.inside)
    print("Coverage Loss:",args.coverage)
    print("Mutex Loss:",args.mutex)
    print("Render Loss:",args.render)

    optimizer = get_optimizer(model,args)
    #scheduler = get_scheduler(optimizer,args)
    loss = get_losses(args)

    if args.pretrain is not None:
        model.load_state_dict(torch.load(args.pretrain))
        print("load pretrained model")

    print("Cuda...")
    _=model.cuda()
    return model,test_loader,optimizer,loss

def get_loss(args,loss,boxes,obj_to_img,inside_box,inside_coords, layout,objs,
    boxes_pred, gene_layout, boxes_refine):
    '''losses'''
    total_loss = 0
    loss_items = {}
    loss_mid = {}
    epoch = 5
    step_weight = [0.1,0.5,1.0]
    for name in loss:
        loss_mid[name] = None
        if name=='box_mse':
            loss_mid[name] = loss[name](boxes_pred,boxes)
        else:
            if epoch>1:
                if name=='gene_ce':
                    loss_mid[name] = step_weight[epoch-2 if epoch<=3 else -1]*loss[name](gene_layout,layout)
                elif name=='mutex':
                    loss_mid[name] = 0.1*loss[name](boxes_pred,obj_to_img,objs)
                    if args.box_refine and args.loss_refine and epoch>2: loss_mid[name]+=loss[name](boxes_refine,obj_to_img,objs)
                elif name=='inside':
                    loss_mid[name] = 0.1*loss[name](boxes_pred,inside_box,obj_to_img)
                    if args.box_refine and args.loss_refine and epoch>2: loss_mid[name]+=loss[name](boxes_refine,inside_box,obj_to_img)
                elif name=='coverage':
                    loss_mid[name] = 0.1*loss[name](boxes_pred,inside_coords,obj_to_img)
                    if args.box_refine and args.loss_refine and epoch>2: loss_mid[name]+=loss[name](boxes_refine,inside_coords,obj_to_img)
                elif name=='render':
                    loss_mid[name] = loss[name](boxes_pred,boxes)
                    if args.box_refine and args.loss_refine and epoch>2: loss_mid[name]+=loss[name](boxes_refine,boxes)
            
            if epoch>2:
                if name=='box_ref_mse':
                    loss_mid[name] = step_weight[epoch-3 if epoch<=4 else -1]*loss[name](boxes_refine,boxes)

        if loss_mid[name] is not None:
            total_loss+=loss_mid[name]
            loss_items[name]=loss_mid[name].item()
    loss_items['total_loss'] = total_loss.item()
    return loss_mid

def draw_graph_with_grad(boxes,inside_box,attrs_with_grad,triples,obj_to_img,triple_to_img,objs,figsize=128):
    NODE_SIZE_RANGE=(200,5000)
    CMAP = get_color_map()
    ROOM_LABELS = get_vocab()["object_idx_to_name"]

    batch_size= obj_to_img.max().item()+1 #N
    boxes_abs = box_rel2abs(boxes,inside_box,obj_to_img) #[N,4]
    
    attrs = attrs_with_grad.detach().cpu() # 144,35
    global_node_indexs=np.arange(attrs.shape[0])
    saliency_nodes, _ = torch.max(attrs_with_grad.grad.data.abs(), dim=1) 
    saliency_nodes = saliency_nodes.detach().cpu() # 144
    graph_drawings=[]
    node_id_start=0
    for i in range(batch_size): #batch中的第i张图片
        index_node = (obj_to_img == i).nonzero(as_tuple=True)[0]
        index_edge = (triple_to_img == i).nonzero(as_tuple=True)[0]
        global_node_indexs_i = [global_node_indexs[i] for i in index_node]
        boxes_i = boxes_abs[index_node]
        node_attrs = attrs[index_node]
        saliency_nodes_i = saliency_nodes[index_node]
        size_nodes_i = (saliency_nodes_i-saliency_nodes_i.min())/(saliency_nodes_i.max()-saliency_nodes_i.min())*(NODE_SIZE_RANGE[1]-NODE_SIZE_RANGE[0])+NODE_SIZE_RANGE[0]
        room_type_i = objs[index_node]
        color_nodes_i = [(CMAP[room_type]/255.).tolist() for room_type in room_type_i]
        edges = triples[index_edge]
        #print(i,node_attrs.shape[0],room_type_i.shape[0],boxes_i.shape[0],edges.shape[0])
        fig = plt.figure(figsize=(figsize,figsize))
        G=nx.Graph()
        pos={}
        for node_id in range(node_attrs.shape[0]):
            #print("\tnode_id",node_id,list(G.nodes))
            global_node_index = global_node_indexs_i[node_id]
            xc,yc,_,_ = boxes_i[node_id].cpu().detach().numpy()
            G.add_node(global_node_index,roomname=ROOM_LABELS[room_type_i[node_id]])
            pos[global_node_index]=[xc,1-yc]
            #print(node_id,room_type,xc,yc)
        node_id_start+=node_attrs.shape[0]
        for edge_id in range(edges.shape[0]):
            
            edge_type = edges[edge_id][1]
            edge_start = edges[edge_id][0].item()
            edge_end = edges[edge_id][2].item()
            #print("\tedge_id",edge_start,edge_end)
            G.add_edge(edge_start,edge_end)
        
        options = {
            "node_color": color_nodes_i,#saliency_nodes_i,
            "node_size":size_nodes_i,
            "edge_color": "#dddddd",
            "width": 4,
            "edge_cmap": plt.cm.Blues,
            "with_labels": False,
        }
        #print(pos)
        #print(list(G.nodes))
        nx.draw(G,pos,**options)
        node_labels = nx.get_node_attributes(G, 'roomname')
        nx.draw_networkx_labels(G, pos, labels=node_labels,font_size=18)
        graph_drawings.append(fig2data(fig))
        fig.clf()
        plt.close()
    return graph_drawings

def draw_boxes(boxes,inside_box,obj_to_img,objs,figsize=128):
    cmap = get_color_map()
    boxes_abs = box_rel2abs(boxes,inside_box,obj_to_img)
    batch_size= obj_to_img.max().item()+1
    imgs = []
    for i in range(batch_size):
        index = (obj_to_img == i).nonzero(as_tuple=True)[0]
        boxes_i = boxes_abs[index]
        room_type_i = objs[index]
        img = np.zeros((figsize,figsize,4), dtype=np.uint8)
        n_rooms = boxes_i.shape[0]
        for j in range(n_rooms):
            drawings = np.zeros((figsize,figsize,4), dtype=np.uint8)
            x,y,w,h = boxes_i[j].cpu().numpy()*figsize
            c = cmap[room_type_i[j]].tolist()
            cv2.rectangle(drawings, (int(x-w/2), int(y-h/2)), (int(x+w/2), int(y+h/2)), (c[0],c[1],c[2],255), -1)
            #c = cmap[-1].tolist()
            cv2.rectangle(drawings, (int(x-w/2), int(y-h/2)), (int(x+w/2), int(y+h/2)), (c[0],c[1],c[2],255),2)
            img+=drawings
        imgs.append(img)
    return imgs

"""def draw_boundary_with_grad(boundary_with_grad):
    boundary_image= boundary_with_grad.detach().cpu() # 20,3,128,128
    boundary_saliency, _ = torch.max(boundary_with_grad.grad.data.abs(), dim=1) 
    boundary_saliency = boundary_saliency.unsqueeze(1).detach().cpu() # 20,1,128,128
    return boundary_image.numpy().transpose(0,2,3,1), boundary_saliency.numpy().transpose(0,2,3,1)
"""
def show_saliency(boundary_image,boundary_saliency,box_imgs,graph_imgs,nrow=20): # 20,3,128,128
    n_plots=4
    batch_size = boundary_image.shape[0]

    fig = plt.figure(figsize=(batch_size//nrow*n_plots*4,4*nrow))
    for i in range(batch_size):
        ax_boxes = fig.add_subplot(nrow,batch_size//nrow*n_plots,n_plots*i+3)
        ax_boxes.imshow(box_imgs[i])
        ax_boxes.axis('off')
        ax_graphs = fig.add_subplot(nrow,batch_size//nrow*n_plots,n_plots*i+4)
        ax_graphs.imshow(graph_imgs[i])
        ax_graphs.axis('off')
        ax_image = fig.add_subplot(nrow,batch_size//nrow*n_plots,n_plots*i+1)
        ax_image.imshow(boundary_image[i])
        ax_image.axis('off')
        ax_saliency = fig.add_subplot(nrow,batch_size//nrow*n_plots,n_plots*i+2)
        ax_saliency.imshow(boundary_saliency[i])
        ax_saliency.axis('off')
        if n_plots*i==0:
            fontsize=20
            ax_boxes.set_title("ground truth layout",fontsize=fontsize)
            ax_graphs.set_title("graph with saliency",fontsize=fontsize)
            ax_image.set_title("input boundary",fontsize=fontsize)
            ax_saliency.set_title("saliency of input boundary",fontsize=fontsize)

    fig.tight_layout()
    return fig